package com.alameda.common.config.mybatis;

import com.alameda.common.aop.CustomDataPermAop;
import com.alameda.common.commonpent.context.SpringContextHolder;
import com.alameda.common.exception.BusinessException;
import com.alameda.common.result.ResultEnum;
import com.alameda.common.utils.PubUtils;
import com.baomidou.mybatisplus.extension.plugins.handler.DataPermissionHandler;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.expression.Alias;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
import net.sf.jsqlparser.parser.CCJSqlParserManager;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.select.FromItem;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.select.Select;
import net.sf.jsqlparser.statement.select.SelectBody;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlSource;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.session.Configuration;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.apache.ibatis.session.SqlSessionFactory;

import java.io.StringReader;
import java.util.List;

/**
 * 数据权限拦截器
 *
 * @author Lv
 * @date 2024/4/9
 */
@Slf4j
@org.springframework.context.annotation.Configuration
@Intercepts({@Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class})})
public class CustomDataPermissionInterceptor implements DataPermissionHandler {

    /**
     * 获取sql
     *
     * @param where             条件
     * @param mappedStatementId 映射语句
     * @return
     */
    @Override
    public Expression getSqlSegment(Expression where, String mappedStatementId) {
        // 获取mybatis sql执行对象
        SqlSessionFactory sqlSessionFactory = SpringContextHolder.getByClass(SqlSessionFactory.class);

        // 获取mybatis sql配置对象
        Configuration configuration = sqlSessionFactory.getConfiguration();
        // 获取sql映射信息
        MappedStatement mappedStatement = getMappedStatement(configuration, mappedStatementId);

        try {
            // 获取本地线程池中的权限隔离sql
            List<String> sqlList = CustomDataPermAop.THREAD_LOCAL.get();
            // 处理拼接sql
            if (sqlList != null && !sqlList.isEmpty()) {
                for (String sql : sqlList) {

                    // 获取查询主表
                    String mainTableName = getMainTableName(mappedStatement);
                    if (PubUtils.isNotEmpty(mainTableName)) {
                        sql = mainTableName + "." + sql;
                    }
                    Expression sqlSegmentExpression = CCJSqlParserUtil.parseCondExpression(sql);

                    if (where == null) {
                        where = sqlSegmentExpression;
                    } else {
                        where = new AndExpression(where, sqlSegmentExpression);
                    }
                }
            }
            return where;
        } catch (JSQLParserException e) {
            log.error("sql解析异常", e);
            throw new BusinessException(ResultEnum.DATA_ERROR, "数据权限异常！");
        }
    }

    /**
     * 获取查询主表
     *
     * @param mappedStatement sql映射信息
     * @return
     */
    private String getMainTableName(MappedStatement mappedStatement) {
        SqlSource sqlSource = mappedStatement.getSqlSource();
        BoundSql boundSql;
        try {
            boundSql = sqlSource.getBoundSql(mappedStatement.getId());
        } catch (Exception exception) {
            return "";
        }
        // 获取执行的SQL语句
        String originalSql = boundSql.getSql();

        CCJSqlParserManager parserManager = new CCJSqlParserManager();
        Select select = null;
        try {
            select = (Select) parserManager.parse(new StringReader(originalSql));
        } catch (JSQLParserException e) {
            e.printStackTrace();
        }
        SelectBody selectBody = select.getSelectBody();

        if (selectBody instanceof PlainSelect) {
            // 向下转型为查询语句后获取主表
            return getMainTable((PlainSelect) selectBody);
        }

        return "";
    }

    /**
     * 获取主表名称
     *
     * @param plainSelect 查询语句对象
     * @return
     */
    private String getMainTable(PlainSelect plainSelect) {
        FromItem fromItem = plainSelect.getFromItem();
        Table table = (Table) fromItem;
        Alias fromItemAlias = table.getAlias();
        String originalTableName = table.getName();
        return fromItemAlias == null ? originalTableName : fromItemAlias.getName();
    }

    /**
     * 获取sql映射信息
     *
     * @param configuration mybatis配置信息
     * @param statementId   sql执行语句
     * @return
     */
    public MappedStatement getMappedStatement(Configuration configuration, String statementId) {
        // 通过 statementId 获取 MappedStatement 对象
        return configuration.getMappedStatement(statementId);
    }
}